/// Memory Manager
use crate::{
    address::*, 
    memory::{
        memory_get_mut, 
        memory_read_bytes, 
        memory_write_bytes
    }, 
    mmu::vpn_to_ppn, 
    page_manager::*, 
    page_table::*, 
    PAGE_SIZE,
};

use std::collections::HashSet;

pub struct MemoryManager {
    root_ppn: PhysPageNum,
    pages: HashSet<PhysPageNum>,
}

impl MemoryManager {
    pub fn new() -> Self {
        let mut mm = MemoryManager {
            root_ppn: alloc_page(),
            pages: HashSet::new(),
        };
        mm.pages.insert(mm.root_ppn);
        mm
    }

    pub fn map(&mut self, vpn: VirtPageNum, ppn: PhysPageNum) {
        let (index0, index1, index2) = vpn.parse();

        let mut pt_ppn = self.root_ppn;
        for (level, index) in [index0, index1, index2].into_iter().enumerate() {
            let paddr = PhysAddr::num_and_offset(pt_ppn, 0);
            let page_table = memory_get_mut::<PageTable>(paddr);
            let pte = &mut page_table[index];
            if pte.is_valid() {
                if level == 2 {
                    panic!("Map a mapped vpn {}", vpn.0);
                }
                pt_ppn = pte.ppn();
            } else {
                pt_ppn = if level == 2 { ppn } else { self.alloc_page() };
                *pte = PageTableEntry::new(pt_ppn, PTEFlags::V);
            }
        }
    }

    pub fn unmap(&mut self, vpn: VirtPageNum) {
        let (index0, index1, index2, _) = VirtAddr::num_and_offset(vpn, 0).parse();
        
        let mut pt_ppn = self.root_ppn;
        let mut need_to_free: Vec<PhysPageNum> = Vec::new();
        for index in [index0, index1, index2].into_iter() {
            let paddr = PhysAddr::num_and_offset(pt_ppn, 0);
            let page_table = memory_get_mut::<PageTable>(paddr);
            let pte = &mut page_table[index];
            if pte.is_valid() {
                pt_ppn = pte.ppn();
                need_to_free.push(pt_ppn);
                *pte = PageTableEntry::new(0.into(), PTEFlags::empty());
            } else {
                break;
            }
        }

        for ppn in need_to_free.into_iter() {
            dealloc_page(ppn);
        }
    }

    pub fn insert(&mut self, vpn: VirtPageNum) -> PhysPageNum {
        let ppn = self.alloc_page();
        self.map(vpn, ppn);
        ppn
    }

    pub fn write_bytes(&mut self, vaddr: VirtAddr, bytes: &[u8]) -> usize {
        match self.divide_vaddr_range_effect(vaddr, bytes.len()) {
            None => 0,
            Some(groups) => {
                let mut write_byte_size: usize = 0;
                for (ppn, begin, end) in groups.into_iter() {
                    let paddr = PhysAddr::num_and_offset(ppn, begin);
                    let byte_size = end - begin;
                    let src = &bytes[write_byte_size..(write_byte_size+byte_size)];
                    write_byte_size += memory_write_bytes(paddr, src);
                }
        
                write_byte_size
            }
        }
    }

    pub fn read_bytes(&mut self, vaddr: VirtAddr, bytes: &mut [u8]) -> usize {
        match self.divide_vaddr_range_effect(vaddr, bytes.len()) {
            None => 0,
            Some(groups) => {
                let mut read_byte_size: usize = 0;
                for (ppn, begin, end) in groups.into_iter() {
                    let paddr = PhysAddr::num_and_offset(ppn, begin);
                    let byte_size = end - begin;
                    let dst = &mut bytes[read_byte_size..(read_byte_size+byte_size)];
                    read_byte_size += memory_read_bytes(paddr, dst);
                }
        
                read_byte_size
            }
        }
    }

    pub fn ppn(&self) -> PhysPageNum {
        self.root_ppn
    }
}

impl MemoryManager {
    fn alloc_page(&mut self) -> PhysPageNum {
        let ppn = alloc_page();
        self.pages.insert(ppn);
        ppn
    }

    fn divide_vaddr_range_effect(&self, vaddr: VirtAddr, size: usize) -> Option<Vec<(PhysPageNum, usize, usize)>> {
        let mut result = Vec::new();

        for (vpn, begin, end) in MemoryManager::divide_vaddr_range(vaddr, size) {
            if let Some(ppn) = vpn_to_ppn(self.root_ppn, vpn) {
                result.push((ppn, begin, end));
            } else {
                return None;
            }
        }

        Some(result)
    }
}

impl MemoryManager {
    fn divide_vaddr_range(vaddr: VirtAddr, size: usize) -> Vec<(VirtPageNum, usize, usize)> {
        let mut result = Vec::new();
        let addr: usize = vaddr.into();

        let mut total_byte_size: usize = 0;
        let mut vpn: usize = addr / PAGE_SIZE;
        let mut begin: usize = addr % PAGE_SIZE;
        let mut end: usize;

        loop {
            if total_byte_size == size {
                break;
            }

            let max_byte_size = PAGE_SIZE - begin;
            let need_byte_size = size - total_byte_size;
            let byte_size = std::cmp::min(max_byte_size, need_byte_size);
            end = byte_size + begin;

            result.push((vpn.into(), begin, end));
            total_byte_size += byte_size;

            vpn += 1;
            begin = 0;
        }

        result
    }
}


impl Drop for MemoryManager {
    fn drop(&mut self) {
        for ppn in &self.pages {
            dealloc_page(ppn.clone());
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{mm::MemoryManager, PAGE_SIZE};

    #[test]
    fn test_divide_vaddr_range() {
        assert_eq!(
            MemoryManager::divide_vaddr_range(23.into(), 2 * PAGE_SIZE),
            vec![(0.into(), 23, PAGE_SIZE), (1.into(), 0, PAGE_SIZE), (2.into(), 0, 23)]
        );

        assert_eq!(
            MemoryManager::divide_vaddr_range((PAGE_SIZE + 343).into(), 10),
            vec![(1.into(), 343, 353)]
        );

        assert_eq!(
            MemoryManager::divide_vaddr_range((13 * PAGE_SIZE + 123).into(),  PAGE_SIZE + 90),
            vec![(13.into(), 123, PAGE_SIZE), (14.into(), 0, 123 + 90)]
        );

        assert_ne!(
            MemoryManager::divide_vaddr_range((13 * PAGE_SIZE + 123).into(),  PAGE_SIZE + 90),
            vec![(10.into(), 123, PAGE_SIZE), (14.into(), 0, 123 + 90)]
        );
    }
}